// The MPC formulation for r^2, m^2 and h, similar to the one in Ni Dong's dissertation.
// One decision variable, T(t_k)
#include "MPC_EW2D_M2_sincos_WithConstraintsTandW.h"
using namespace Ipopt;

const double pi  = 3.141592653;
const double k_w = 1.678251e-12;
const double a_w = 6.194597e-01;
const double E_w = 2.835225e-03;		// Unit = EV
const double k_c = 4.399286e-17;
const double a_c = 5.604130e-01;
const double E_c = 1.624391e-03;		// Unit = EV
const double T0_W0 = -6.581760e-01;
const double T0_W1 = 5.178483e-01;
const double T1_W0 = 3.105538e-03;
const double T1_W1 = -2.409770e-03;
const double T2_W0 = -4.879777e-06;
const double T2_W1 = 3.741534e-06;
const double T3_W0 = 2.553569e-09;
const double T3_W1 = -1.937348e-09;
const double k_B = 8.617343e-5; // Boltzmann's constant, Unit = ev/K

// -----------------------------------------------------------------------
/* Constructor. */
MPC_Formulation::MPC_Formulation()
{
	int m,n;
	h_set = 100.0;
	r2_set = 15;
	m2_set = 0.2;

	t_end = 100.0;
	dt    = 5.0;
	mode  = 50;
	LatticeSize = 100;
	q_h   = 0.0;
	q_r   = 0.0/(r2_set*r2_set);
//	q_m   = 0.0;
//	q_r   = 0.0;
	q_m   = 1.0/(m2_set*m2_set);

	cov_z1 = new double[(1+mode)*(1+mode)];
	cov_z2 = new double[(1+mode)*(1+mode)];
	cov_z3 = new double[(1+mode)*(1+mode)];
	cov_z4 = new double[(1+mode)*(1+mode)];

	Km1 = new double*[mode+1];
	Km2 = new double*[mode+1];
	Km3 = new double*[mode+1];
	Km4 = new double*[mode+1];

	double dx = pi/LatticeSize;
	for(m=0;m<=mode;m++){
		Km1[m] = new double[mode+1];
		Km2[m] = new double[mode+1];
		Km3[m] = new double[mode+1];
		Km4[m] = new double[mode+1];

		for(n=0;n<=mode;n++){
			if(m==0||2*m==LatticeSize||n==0||2*n==LatticeSize){
				Km1[m][n] = 0.0;
			}else{
				Km1[m][n] = 4.0/(pi*pi)*sin(pi*m/LatticeSize)*sin(pi*m/LatticeSize);
			}

			if(m==0){
				Km2[m][n] = 0.0;
			}else{
				Km2[m][n] = 4.0/(pi*pi)*sin(pi*m/LatticeSize)*sin(pi*m/LatticeSize);
			}

			if(m==0||2*m==LatticeSize){
				Km3[m][n] = 0.0;
			}else{
				Km3[m][n] = 4.0/(pi*pi)*sin(pi*m/LatticeSize)*sin(pi*m/LatticeSize);
			}

			if(n==0 || 2*n==LatticeSize || m==0){
				Km4[m][n] = 0.0;
			}else{
				Km4[m][n] = 4.0/(pi*pi)*sin(pi*m/LatticeSize)*sin(pi*m/LatticeSize);
			}
		}
	}
	reset();
}
// -----------------------------------------------------------------------
MPC_Formulation::~MPC_Formulation()
{
	delete []cov_z1;
	delete []cov_z2;
	delete []cov_z3;
	delete []cov_z4;
	int m;
	for(m=0;m<=mode;m++){
		delete [] Km1[m];
		delete [] Km2[m];
		delete [] Km3[m];
		delete [] Km4[m];
	}
	delete [] Km1;
	delete [] Km2;
	delete [] Km3;
	delete [] Km4;
}
// -----------------------------------------------------------------------
void MPC_Formulation::parameterize(double CurrentTime, double** inputs)
{
	int    m,n;
	int    index;
	for(m=0;m<=mode;m++){
		for(n=0;n<=mode;n++){
			index = index2D(m,n);
			cov_z1[index] = inputs[0][index];
			cov_z2[index] = inputs[1][index];
			cov_z3[index] = inputs[2][index];
			cov_z4[index] = inputs[3][index];
		}
	}
	time = CurrentTime;
}
// -----------------------------------------------------------------------
void MPC_Formulation::reset(){
	T_old = 610;
	W_old = 1.0;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_nlp_info(Index& n, Index& m, Index& nnz_jac_g,
                         Index& nnz_h_lag, IndexStyleEnum& index_style)
{
  // The problem described in MPC_Formulation.hpp has 2 variables, x1 = T, & x2 = W
  n = 2;

  // one inequality constraint: C > C_{min}
  m = 1;

  // 2 nonzeros in the jacobian (one for x1, and one for x2),
  nnz_jac_g = 2;

  // and 2 nonzeros in the hessian of the lagrangian
  // (one in the hessian of the objective for x2,
  //  and one in the hessian of the constraints for x1)
  nnz_h_lag = 4;

  // We use the standard fortran index style for row/col entries
  index_style = C_STYLE;

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_bounds_info(Index n, Number* x_l, Number* x_u,
                            Index m, Number* g_l, Number* g_u)
{
  // here, the n and m we gave IPOPT in get_nlp_info are passed back to us.
  // If desired, we could assert to make sure they are what we think they are.
  assert(n == 2);
  assert(m == 1);

  // x1 has a lower bound of -1 and an upper bound of 1
  x_l[0] = max(600.0,T_old-25.0);
  x_u[0] = min(750.0,T_old+25.0);

  // x2 has no upper or lower bound, so we set them to
  // a large negative and a large positive number.
  // The value that is interpretted as -/+infinity can be
  // set in the options, but it defaults to -/+1e19
  x_l[1] = max(0.1,W_old-0.25);
  x_u[1] = min(1.0,W_old+0.25);

  // we have one equality constraint, so we set the bounds on this constraint
  // to be equal (and zero).
  g_l[0] = 0.1*h_set/t_end;
  g_u[0] = 2e19;

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_starting_point(Index n, bool init_x, Number* x,
                               bool init_z, Number* z_L, Number* z_U,
                               Index m, bool init_lambda,
                               Number* lambda)
{
  // Here, we assume we only have starting values for x, if you code
  // your own NLP, you can provide starting values for the others if
  // you wish.
  assert(init_x == true);
  assert(init_z == false);
  assert(init_lambda == false);

  // we initialize x in bounds, in the upper right quadrant
  x[0] = T_old;
  x[1] = W_old;
  //x[0] = 610;
  //x[1] = 0.5;

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_f(Index n, const Number* x, bool new_x, Number& obj_value)
{
  // return the value of the objective function
  // J = q_h*(h_set-h_f)^2+q_r*(r2_s-r2_f)^2

  Number T = x[0];
  Number W = x[1];

  double h_f;
  double r2_f;
  double m2_f;
  
  double C;
  double C2;
  double sigma2;

  C = eq_C(T,W);
  C2 = eq_C2(T,W);
  sigma2 = eq_sigma2(T,W);

  h_f = eq_hf(C);
  r2_f = eq_r2_f(C2,sigma2);
  m2_f = eq_m2_f(C2,sigma2);
  
  obj_value = q_h*(h_set-h_f)*(h_set-h_f)+q_r*(r2_set-r2_f)*(r2_set-r2_f)+q_m*(m2_set-m2_f)*(m2_set-m2_f);

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_grad_f(Index n, const Number* x, bool new_x, Number* grad_f)
{
  // return the gradient of the objective function grad_{x} f(x)

  double T = x[0];
  double W = x[1];

  double C = eq_C(T,W);
  double C2 = eq_C2(T,W);
  double sigma2 = eq_sigma2(T,W);
  double hf = eq_hf(C);
  double r2_f = eq_r2_f(C2,sigma2);
  double m2_f = eq_m2_f(C2,sigma2);
  
  double dhfdC = eq_dhfdC(C);
  double dr2_fdC2 = eq_dr2_fdC2(C2,sigma2);
  double dr2_fdsigma2 = eq_dr2_fdsigma2(C2,sigma2);
  double dm2_fdC2 = eq_dm2_fdC2(C2,sigma2);
  double dm2_fdsigma2 = eq_dm2_fdsigma2(C2,sigma2);

  double dCdT  = eq_dCdT(T,W);
  double dCdW = eq_dCdW(T,W);
  double dC2dT = eq_dC2dT(T,W);
  double dC2dW = eq_dC2dW(T,W);
  double dsigma2dT = eq_dsigma2dT(T,W);
  double dsigma2dW = eq_dsigma2dW(T,W);

  // grad_{x1} f(x): x1 is not in the objective
  grad_f[0] = -2*q_h*(h_set-hf)*dhfdC*dCdT-2*q_r*(r2_set-r2_f)*(dr2_fdC2*dC2dT+dr2_fdsigma2*dsigma2dT)-2*q_m*(m2_set-m2_f)*(dm2_fdC2*dC2dT+dm2_fdsigma2*dsigma2dT);

  // grad_{x2} f(x):
  grad_f[1] = -2*q_h*(h_set-hf)*dhfdC*dCdW-2*q_r*(r2_set-r2_f)*(dr2_fdC2*dC2dW+dr2_fdsigma2*dsigma2dW)-2*q_m*(m2_set-m2_f)*(dm2_fdC2*dC2dW+dm2_fdsigma2*dsigma2dW);

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_g(Index n, const Number* x, bool new_x, Index m, Number* g){
  // return the value of the constraints: g(x)
  Number T = x[0];
  Number W = x[1];

  g[0] = eq_C(T,W);

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_jac_g(Index n, const Number* x, bool new_x,
                       Index m, Index nele_jac, Index* iRow, Index *jCol,
                       Number* values)
{
  if (values == NULL) {
    // return the structure of the jacobian of the constraints

    // element at 1,1: grad_{x1} g_{1}(x)
    iRow[0] = 0;    jCol[0] = 0;

    // element at 1,2: grad_{x2} g_{1}(x)
    iRow[1] = 0;    jCol[1] = 1;
  }
  else {
    // return the values of the jacobian of the constraints
	double T = x[0];
	double W = x[1];

    // element at 1,1: grad_{x1} g_{1}(x)
    values[0] = eq_dCdT(T,W);

    // element at 1,2: grad_{x1} g_{1}(x)
    values[1] = eq_dCdW(T,W);
  }

  return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_h(Index n, const Number* x, bool new_x,
                   Number obj_factor, Index m, const Number* lambda,
                   bool new_lambda, Index nele_hess, Index* iRow,
                   Index* jCol, Number* values)
{
  if (values == NULL) {
    // return the structure. This is a symmetric matrix, fill the lower left
    // triangle only.
    iRow[0] = 0;    jCol[0] = 0;
    iRow[1] = 1;    jCol[1] = 0;
	iRow[2] = 1;    jCol[2] = 1;
  }
  else {
    // return the values
    // fill the objective portion
    values[0] = 0;
    values[1] = 0;
	values[2] = 0;
	// add the portion for the first constraint
	values[0] += 0;
	values[1] += 0;
	values[2] += 0;
  }
  return false; // USE Quasi-Newton method
}
// -----------------------------------------------------------------------
void MPC_Formulation::finalize_solution(SolverReturn status,
                              Index n, const Number* x, const Number* z_L, const Number* z_U,
                              Index m, const Number* g, const Number* lambda,
                              Number obj_value,
			      const IpoptData* ip_data,
			      IpoptCalculatedQuantities* ip_cq)
{
  // here is where we would store the solution to variables, or write to a file, etc
  // so we could use the solution. Since the solution is displayed to the console,
  // we currently do nothing here.

  // Save the current solution, which will be used as the initial solution.

	// If time == t_end, the result is discard, keep the laatest result.
	if(time!=t_end){
		T_old = x[0];
		W_old = x[1];
	}

	double C;
	double C2;
	double sigma2;

    double h_f;
    double r2_f;
	double m2_f;

    C = eq_C(T_old,W_old);
    C2 = eq_C2(T_old,W_old);
    sigma2 = eq_sigma2(T_old,W_old);

    h_f = eq_hf(C);
    r2_f = eq_r2_f(C2,sigma2);
	m2_f = eq_m2_f(C2,sigma2);

	printf("*** time = %f, T = %f, W = %f,\nr_f = %f, h_f = %f, m2_f = %f\n",time,T_old,W_old,r2_f,h_f,m2_f);
}



// -----------------------------------------------------------------------
double MPC_Formulation::eq_C(double T,double W)
{
	return W*(1-k_w*pow(W,-a_w)*exp(k_B*T/E_w));
}

double MPC_Formulation::eq_dCdT(double T,double W)
{
	return -W*k_w*exp(k_B*T/E_w)*k_B/(pow(W,a_w)*E_w);
}

double MPC_Formulation::eq_dCdW(double T,double W)
{
	return 1-k_w*exp(k_B*T/E_w)*(1-a_w)*pow(W,-a_w);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_C2(double T,double W)
{
	return k_c*exp(k_B*T/E_c)*pow(W,-a_c)/(LatticeSize*LatticeSize);
}

double MPC_Formulation::eq_dC2dT(double T,double W)
{
	return k_c/(LatticeSize*LatticeSize*pow(W,a_c))*exp(k_B*T/E_c)*k_B/E_c;
}

double MPC_Formulation::eq_dC2dW(double T,double W)
{
	return k_c/(LatticeSize*LatticeSize)*exp(k_B*T/E_c)*(-a_c)*pow(W,-a_c-1);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_sigma2(double T,double W)
{
//	double index = (a_t+k_t*W)*T-a_v-k_v*W;
//	return pi*pi/(LatticeSize*LatticeSize)*(1+exp(index));
	return (T0_W0+T0_W1*W)+(T1_W0+T1_W1*W)*T+(T2_W0+T2_W1*W)*T*T+(T3_W0+T3_W1*W)*T*T*T;
}

double MPC_Formulation::eq_dsigma2dT(double T,double W)
{
//	return pi*pi/(LatticeSize*LatticeSize)*exp((a_t+k_t*W)*T-a_v-k_v*W)*(a_t+k_t*W);
	return (T1_W0+T1_W1*W)+2*(T2_W0+T2_W1*W)*T+3*(T3_W0+T3_W1*W)*T*T;
}

double MPC_Formulation::eq_dsigma2dW(double T,double W)
{
//	return 	pi*pi/(LatticeSize*LatticeSize)*exp((a_t+k_t*W)*T-a_v-k_v*W)*(k_t*T-k_v);
	return T0_W1+T1_W1*T+T2_W1*T*T+T3_W1*T*T*T;
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_K1(double C2,int m, int n,double delta_t)
{
	return exp(-8*(m*m+n*n)*C2*delta_t);
}

double MPC_Formulation::eq_dK1dC2(double C2,int m,int n,double delta_t)
{
	return eq_K1(C2,m,n,delta_t)*(-8*(m*m+n*n)*delta_t);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_K2(double C2,int m,int n,double delta_t)
{
	return (exp(-8*(m*m+n*n)*C2*delta_t)-1)/(-8*(m*m+n*n)*C2);
}

double MPC_Formulation::eq_dK2dC2(double C2,int m,int n,double delta_t)
{
	double K = -8*(m*m+n*n);
	double temp1 = exp(K*delta_t*C2);
	return (temp1*K*delta_t*K*C2-K*(temp1-1))/(K*K*C2*C2);
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_r2_f(double C2,double sigma2)
{
	int m,n;

	double r2_f = 0.0;
	double temp1;
	int    index;
	double lambda;
	for(m=0;m<=mode;m++){
		for(n=0;n<=mode;n++){
			if(m!=0||n!=0){
				lambda = -4.0*C2*(m*m+n*n);
				temp1 = exp(2*lambda*(t_end-time));
				index = index2D(m,n);
				r2_f += (cov_z1[index]+cov_z2[index]+cov_z3[index]+cov_z4[index])*temp1+4*sigma2*(temp1-1)/(2*lambda);
			}
		}
	}
	return r2_f/=(pi*pi);
}

double MPC_Formulation::eq_dr2_fdC2(double C2,double sigma2)
{
	int m,n;
	double dr2_fdC2 = 0.0;
	double lambda;
	double temp1;
	int    index;
	for(m=0;m<=mode;m++){
		for(n=0;n<=mode;n++){
			if(m!=0||n!=0){
				lambda = -4.0*C2*(m*m+n*n);
				temp1 = exp(2*lambda*(t_end-time));
				index = index2D(m,n);
				dr2_fdC2 += ((cov_z1[index]+cov_z2[index]+cov_z3[index]+cov_z4[index])*temp1*2*(t_end-time)
					        +sigma2/(lambda*lambda)*(4*temp1*(t_end-time)*lambda-2*(temp1-1)))*(-4*(m*m+n*n));
			}
		}
	}
	dr2_fdC2 /= (pi*pi);
	return dr2_fdC2;
}

double MPC_Formulation::eq_dr2_fdsigma2(double C2,double sigma2)
{
	int m,n;
	double dr2_fdsigma2 = 0.0;
	for(m=0;m<=mode;m++){
		for(n=0;n<=mode;n++){
			if(m!=0||n!=0){
				dr2_fdsigma2 += 4*eq_K2(C2,m,n,t_end-time);
			}
		}
	}
	dr2_fdsigma2 /= (pi*pi);
	return dr2_fdsigma2;
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_m2_f(double C2,double sigma2){
	int m,n;
	double m2_f = 0.0;
	double temp1;
	double lambda;
	int    index;
	for(m=0;m<=mode;m++){
		for(n=0;n<=mode;n++){
			if(m!=0||n!=0){
				lambda = -4.0*C2*(m*m+n*n);
				temp1 = exp(2*lambda*(t_end-time));
				index = index2D(m,n);
				m2_f += (Km1[m][n]*cov_z1[index]+Km2[m][n]*cov_z2[index]+Km3[m][n]*cov_z3[index]+Km4[m][n]*cov_z4[index])*temp1
					    +(Km1[m][n]+Km2[m][n]+Km3[m][n]+Km4[m][n])*sigma2*(temp1-1)/(2*lambda);
			}
		}
	}
	return m2_f;
}

double MPC_Formulation::eq_dm2_fdC2(double C2,double sigma2){
	int m,n;
	double dm2_fdC2 = 0.0;
	double lambda;
	double temp1;
	int    index;
	for(m=0;m<=mode;m++){
		for(n=0;n<=mode;n++){
			if(m!=0||n!=0){
				lambda = -4.0*C2*(m*m+n*n);
				temp1 = exp(2*lambda*(t_end-time));
				index = index2D(m,n);
				dm2_fdC2 += ((cov_z1[index]*Km1[m][n]+cov_z2[index]*Km2[m][n]+cov_z3[index]*Km3[m][n]+cov_z4[index]*Km4[m][n])*temp1*2*(t_end-time)
					       +sigma2/(4*lambda*lambda)*(Km1[m][n]+Km2[m][n]+Km3[m][n]+Km4[m][n])*(4*temp1*(t_end-time)*lambda-2*(temp1-1)))*(-4*(m*m+n*n));
			}
		}
	}
	return dm2_fdC2;
}

double MPC_Formulation::eq_dm2_fdsigma2(double C2,double sigma2){
	int m,n;
	double dm2_fdsigma2 = 0.0;
	for(m=0;m<=mode;m++){
		for(n=0;n<=mode;n++){
			if(m!=0||n!=0){
				dm2_fdsigma2 += (Km1[m][n]+Km2[m][n]+Km3[m][n]+Km4[m][n])*eq_K2(C2,m,n,t_end-time);
			}
		}
	}
	return dm2_fdsigma2;
}
// -----------------------------------------------------------------------
double MPC_Formulation::eq_hf(double C)
{
	return C*(t_end-time);
}

double MPC_Formulation::eq_dhfdC(double C)
{
	return t_end-time;
}
// -----------------------------------------------------------------------
void MPC_Formulation::get_solution(double *x)
{
	x[0] = T_old;
	x[1] = W_old;
}